import os
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
import albumentations as A
import cv2
import SimpleITK as sitk
from config import zhongShanParameter as para



class TrainDataset(data.Dataset):
    """
    dataloader for skin lesion segmentation tasks
    """
    def __init__(self, root, id_path,channel=1):
        self.root = root
        self.channel=channel
        with open(id_path, 'r') as f:
            self.ids = f.read().splitlines()

        self.size = len(self.ids)

        self.img_transform = transforms.Compose([
            transforms.ToTensor(),#将所有数除以255，将数据归一化到【0，1】,但是这里数值范围并没有改变？
            # transforms.Normalize([0.485],#减去均值，再除以方差，这里应该计算自己数据集的参数
            #                      [0.229])
        ])
        self.seg_transform = transforms.Compose([
            transforms.ToTensor()])#将所有数除以255，将数据归一化到【0，1】
        
        self.transform = A.OneOf(
            [
                A.ShiftScaleRotate(shift_limit=0.15, scale_limit=0.15, rotate_limit=25, p=0.5, border_mode=0),
                #A.ColorJitter(),
                A.HorizontalFlip(),
                A.VerticalFlip(),
                #下面是新加的数据增强,注意BST-ST论文没用ShiftScaleRotate
                A.Transpose(),
                A.GaussNoise(),
                A.ZoomBlur()

            ]
        )

    def __getitem__(self, index):
        id = self.ids[index]
        ct_path = os.path.join(self.root, id.split(' ')[0])
        seg_path = os.path.join(self.root, id.split(' ')[1])

        image = sitk.ReadImage(ct_path,sitk.sitkFloat32)#为了和权重数据类型一致，要设置float32
        seg = sitk.ReadImage(seg_path, sitk.sitkFloat32)

        image = sitk.GetArrayFromImage(image)
        seg = sitk.GetArrayFromImage(seg)
        # image = mynormalize(image)
        # seg = seg/255#待改正
        image=np.resize(image,(image.shape[0],image.shape[0],1))
        seg = np.resize(seg,(image.shape[0],image.shape[0],1))

        transformed = self.transform(image=image, mask=seg)
        image = self.img_transform(transformed['image'])


        if self.channel==3:# 通道复制，因为预训练模型是三通道的
            image = image.repeat(3,  1, 1)
        seg = self.seg_transform(transformed['mask'])

        return image, seg

    def __len__(self):
        return self.size





class TestDataset(data.Dataset):
    """

    """

    def __init__(self, root, id_path,channel):
        self.root = root
        self.channel=channel
        with open(id_path, 'r') as f:
            self.ids = f.read().splitlines()

        self.size = len(self.ids)

        self.img_transform = transforms.Compose([
            transforms.ToTensor(),
            # transforms.Normalize([0.485],
            #                      [0.229])
        ])
        self.seg_transform = transforms.Compose([
            transforms.ToTensor()])

    def __getitem__(self, index):
        id = self.ids[index]
        ct_path = os.path.join(self.root, id.split(' ')[0])
        seg_path = os.path.join(self.root, id.split(' ')[1])

        image = sitk.ReadImage(ct_path, sitk.sitkFloat32)
        seg = sitk.ReadImage(seg_path, sitk.sitkFloat32)

        image = sitk.GetArrayFromImage(image)
        seg = sitk.GetArrayFromImage(seg)
        # image=mynormalize(image)
        # seg = seg/255#待改正
        image = np.resize(image, (image.shape[0], image.shape[0], 1))
        seg = np.resize(seg, (image.shape[0], image.shape[0], 1))

        image = self.img_transform(image)
        if self.channel == 3:  # 通道复制，因为预训练模型是三通道的
            image = image.repeat(3, 1, 1)
        seg = self.seg_transform(seg)
        return image, seg

    def __len__(self):
        return self.size


def getTrainLoader( root, id_path, batchsize, shuffle=True, num_workers=4, pin_memory=True,channel=1):

    dataset = TrainDataset( root, id_path,
                                  channel=channel)
    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=batchsize,
                                  shuffle=shuffle,
                                  num_workers=num_workers,
                                  pin_memory=pin_memory)
    return data_loader
def getTestLoader( root, id_path, batchsize, shuffle=False, num_workers=4, pin_memory=True,channel=1):

    dataset = TestDataset( root, id_path,
                                  channel=channel)
    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=batchsize,
                                  shuffle=shuffle,
                                  num_workers=num_workers,
                                  pin_memory=pin_memory)
    return data_loader

if __name__ == '__main__':
    dl=getTrainLoader(para.cut2d_save_path,para.train2d_choose_id_path,batchsize = 4)
    size=448

    for image,tg in dl:
        image=torch.tensor(image)
        print(torch.max(image),torch.min(image))



    # path = 'data/'
    # tt = SkinDataset(path + 'data_train.npy', path + 'mask_train.npy')
    # for i in range(50):
    #     img, gt = tt.__getitem__(i)
    #
    #     img = torch.transpose(img, 0, 1)
    #     img = torch.transpose(img, 1, 2)
    #     img = img.numpy()
    #     gt = gt.numpy()
    #
    #     plt.imshow(img)
    #     plt.savefig('vis/'+str(i)+".jpg")
    #
    #     plt.imshow(gt[0])
    #     plt.savefig('vis/'+str(i)+'_gt.jpg')
